"""Evaluates the model"""

import argparse
import logging
import os

import numpy as np
import torch
import yaml
from torch.autograd import Variable
from tqdm import tqdm

import utils
import model.mobilenetv3 as net
import data.data_loader_224 as data_loader

parser = argparse.ArgumentParser()
parser.add_argument('--model_dir', default='experiments/mobile_net',
                    help="Directory containing config.yaml")


def evaluate(model, loss_fn, dataloader, metrics, params):
    """Evaluate the model on `num_steps` batches.

    Args:
        model: (torch.nn.Module) the neural network
        loss_fn: a function that takes batch_output and batch_labels and computes the loss for the batch
        dataloader: (DataLoader) a torch.utils.data.DataLoader object that fetches data
        metrics: (dict) a dictionary of functions that compute a metric using the output and labels of each batch
        params: (Params) hyperparameters
        num_steps: (int) number of batches to train on, each of size params.batch_size
    """

    # set model to evaluation mode
    model.eval()

    # summary for current eval loop
    summ = []
    loss_avg = utils.RunningAverage()

    with tqdm(total=len(dataloader), desc='Testing:') as t:
        # compute metrics over the dataset
        for data_batch, labels_batch in dataloader:

            # move to GPU if available
            if config['device']['gpu'] and torch.cuda.is_available():
                data_batch, labels_batch = data_batch.cuda(non_blocking=True), labels_batch.cuda(non_blocking=True)
            # convert to torch Variables
            data_batch, labels_batch = Variable(data_batch), Variable(labels_batch)

            # compute model output
            output_batch = model(data_batch)
            loss = loss_fn(output_batch, labels_batch)

            # extract data from torch Variable, move to cpu, convert to numpy arrays
            output_batch = output_batch.data.cpu().numpy()
            labels_batch = labels_batch.data.cpu().numpy()

            # compute all metrics on this batch
            summary_batch = {metric: metrics[metric](output_batch, labels_batch)
                             for metric in metrics}
            summary_batch['loss'] = loss.item()
            summ.append(summary_batch)

            # update the average loss
            loss_avg.update(loss.item())

            t.set_postfix(loss='{:05.3f}'.format(loss_avg()))
            t.update()

    # compute mean of all metrics in summary
    metrics_mean = {metric: np.mean([x[metric]
                                     for x in summ]) for metric in summ[0]}
    metrics_string = " ; ".join("{}: {:05.3f}".format(k, v)
                                for k, v in metrics_mean.items())
    logging.info("- Eval metrics : " + metrics_string)
    return metrics_mean


if __name__ == '__main__':
    """
        Evaluate the model on the test set.
    """
    # Load the parameters
    args = parser.parse_args()

    # Get the logger
    utils.set_logger(os.path.join(args.model_dir, 'evaluate.log'))

    yaml_path = os.path.join(args.model_dir, 'config.yaml')
    # print(yaml_path)
    with open(yaml_path, "r", encoding='utf-8') as f:
        config = yaml.safe_load(f)
    assert os.path.isfile(
        yaml_path), "No yaml configuration file found at {}".format(yaml_path)

    # use GPU if available
    device = torch.device(
        f"cuda:{config['device']['num']}" if torch.cuda.is_available() and config['device']['gpu'] else "cpu")
    logging.info(f'Using device: {device}')
    # Set the random seed for reproducible experiments
    torch.manual_seed(42)
    if config['device']['gpu']:
        torch.cuda.manual_seed(42)


    # Create the input data pipeline
    logging.info("Creating the dataset...")

    # fetch dataloaders
    dataloaders = data_loader.fetch_dataloader(['test'], config)
    test_dl = dataloaders['test']

    logging.info("- done.")

    # Define the model
    model = net.MobileNetV3(config).to(device)

    loss_fn = net.loss_fn
    metrics = net.metrics

    logging.info("Starting evaluation")

    best_checkpoint = os.path.join(args.model_dir, 'best.pth.tar')
    restore_file = 'best' if os.path.exists(best_checkpoint) else None
    # Reload weights from the saved file
    utils.load_checkpoint(os.path.join(
        args.model_dir, restore_file + '.pth.tar'), model)

    # Evaluate
    test_metrics = evaluate(model, loss_fn, test_dl, metrics, config)
    save_path = os.path.join(
        args.model_dir, "metrics_test_{}.json".format(restore_file))
    utils.save_dict_to_json(test_metrics, save_path)
